> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/mlfoundations/open_clip/llms.txt
> Use this file to discover all available pages before exploring further.

# CoCaLoss

> Combined contrastive and captioning loss for training CoCa (Contrastive Captioner) models

## Overview

The `CoCaLoss` class extends `ClipLoss` to support training CoCa models, which combine contrastive learning with autoregressive caption generation. It computes a weighted combination of:

1. **Contrastive loss** - CLIP-style image-text matching (inherited from ClipLoss)
2. **Caption loss** - Cross-entropy loss for next-token prediction in caption generation

This dual-objective training enables models to both match images with text descriptions and generate natural language captions.

## Class Definition

```python theme={null}
from open_clip import CoCaLoss
```

## Initialization Parameters

<ParamField path="caption_loss_weight" type="float" required>
  Weight applied to the caption generation loss. Controls the relative importance of caption quality vs. contrastive matching.
</ParamField>

<ParamField path="clip_loss_weight" type="float" required>
  Weight applied to the contrastive loss. Set to 0 to train only the caption decoder.
</ParamField>

<ParamField path="pad_id" type="int" default="0">
  Padding token ID. Positions with this token are ignored when computing caption loss.
</ParamField>

<ParamField path="local_loss" type="bool" default="False">
  If True, computes contrastive loss only between local and gathered features. See ClipLoss documentation for details.
</ParamField>

<ParamField path="gather_with_grad" type="bool" default="False">
  If True, gathers features with gradient flow enabled for contrastive loss.
</ParamField>

<ParamField path="cache_labels" type="bool" default="False">
  If True, caches ground truth labels for contrastive loss.
</ParamField>

<ParamField path="rank" type="int" default="0">
  Current process rank in distributed training.
</ParamField>

<ParamField path="world_size" type="int" default="1">
  Total number of processes in distributed training.
</ParamField>

<ParamField path="use_horovod" type="bool" default="False">
  If True, uses Horovod for distributed operations instead of torch.distributed.
</ParamField>

## Attributes

Inherits all attributes from `ClipLoss`, plus:

* **clip\_loss\_weight**: Weight for contrastive loss component
* **caption\_loss\_weight**: Weight for caption generation loss component
* **caption\_loss**: CrossEntropyLoss module with `ignore_index=pad_id`

## Key Methods

### forward

```python theme={null}
def forward(
    self,
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    logits: torch.Tensor,
    labels: torch.Tensor,
    logit_scale: torch.Tensor,
    output_dict: bool = False,
) -> Union[Tuple[torch.Tensor, torch.Tensor], Dict[str, torch.Tensor]]:
```

Computes the combined CoCa loss.

**Parameters:**

* `image_features`: Normalized contrastive image features of shape `(batch_size, embed_dim)`
* `text_features`: Normalized contrastive text features of shape `(batch_size, embed_dim)`
* `logits`: Caption generation logits of shape `(batch_size, seq_len, vocab_size)`
* `labels`: Target token IDs for caption generation of shape `(batch_size, seq_len)`
* `logit_scale`: Temperature parameter for contrastive loss (typically `model.logit_scale.exp()`)
* `output_dict`: If True, returns dict with named losses, else returns tuple

**Returns:**

* If `output_dict=False`: Tuple of `(clip_loss, caption_loss)` - both weighted by their respective coefficients
* If `output_dict=True`: Dictionary with keys `"contrastive_loss"` and `"caption_loss"`

## Usage Example

```python theme={null}
import torch
from open_clip import create_model, CoCaLoss

# Create CoCa model
model = create_model('coca_ViT-B-32', pretrained=False)

# Create loss function
loss_fn = CoCaLoss(
    caption_loss_weight=2.0,  # Caption loss is weighted 2x
    clip_loss_weight=1.0,     # Standard contrastive loss weight
    pad_id=0
)

# Training loop
images = torch.randn(16, 3, 224, 224)
captions = torch.randint(0, 49408, (16, 77))

# Forward pass through model
output = model(images, captions)

# Extract required tensors
image_features = output['image_features']  # Contrastive features
text_features = output['text_features']    # Contrastive features
logits = output['logits']                  # Caption logits
labels = output['labels']                  # Shifted caption targets
logit_scale = output['logit_scale']

# Compute loss
clip_loss, caption_loss = loss_fn(
    image_features,
    text_features,
    logits,
    labels,
    logit_scale
)

total_loss = clip_loss + caption_loss
total_loss.backward()
```

## Distributed Training Example

```python theme={null}
import torch.distributed as dist
from open_clip import create_model, CoCaLoss

# Initialize distributed
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()

# Create model and loss
model = create_model('coca_ViT-L-14').to(rank)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

loss_fn = CoCaLoss(
    caption_loss_weight=2.0,
    clip_loss_weight=1.0,
    pad_id=0,
    rank=rank,
    world_size=world_size,
    cache_labels=True
)

for images, captions in dataloader:
    images = images.to(rank)
    captions = captions.to(rank)
    
    output = model(images, captions)
    
    clip_loss, caption_loss = loss_fn(
        output['image_features'],
        output['text_features'],
        output['logits'],
        output['labels'],
        output['logit_scale']
    )
    
    total_loss = clip_loss + caption_loss
    total_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

## Dictionary Output for Logging

```python theme={null}
loss_fn = CoCaLoss(
    caption_loss_weight=2.0,
    clip_loss_weight=1.0
)

output = model(images, captions)

loss_dict = loss_fn(
    output['image_features'],
    output['text_features'],
    output['logits'],
    output['labels'],
    output['logit_scale'],
    output_dict=True
)

print(loss_dict)
# {'contrastive_loss': tensor(4.2), 'caption_loss': tensor(3.1)}

# Easy logging to wandb/tensorboard
for key, value in loss_dict.items():
    logger.log({key: value.item()})
```

## Caption-Only Training

```python theme={null}
# Train only the caption decoder (freeze contrastive learning)
loss_fn = CoCaLoss(
    caption_loss_weight=1.0,
    clip_loss_weight=0.0,  # Disable contrastive loss
    pad_id=0
)

# Clip loss will be zero tensor (no backward pass)
clip_loss, caption_loss = loss_fn(...)
print(clip_loss)  # tensor(0.)
```

## Custom Loss Weighting Strategy

```python theme={null}
class AdaptiveCoCaLoss(CoCaLoss):
    """Dynamically adjust loss weights during training."""
    
    def __init__(self, *args, initial_caption_weight=2.0, **kwargs):
        super().__init__(
            caption_loss_weight=initial_caption_weight,
            *args,
            **kwargs
        )
        self.step = 0
    
    def forward(self, *args, **kwargs):
        # Gradually increase caption loss weight
        self.caption_loss_weight = 2.0 + (self.step / 10000)
        self.step += 1
        return super().forward(*args, **kwargs)

loss_fn = AdaptiveCoCaLoss(
    clip_loss_weight=1.0,
    pad_id=0
)
```

## Monitoring Loss Components

```python theme={null}
import torch
from collections import defaultdict

loss_fn = CoCaLoss(
    caption_loss_weight=2.0,
    clip_loss_weight=1.0
)

metrics = defaultdict(list)

for epoch in range(num_epochs):
    for images, captions in dataloader:
        output = model(images, captions)
        
        clip_loss, caption_loss = loss_fn(
            output['image_features'],
            output['text_features'],
            output['logits'],
            output['labels'],
            output['logit_scale']
        )
        
        # Track both components
        metrics['clip_loss'].append(clip_loss.item())
        metrics['caption_loss'].append(caption_loss.item())
        metrics['total_loss'].append((clip_loss + caption_loss).item())
        
        total_loss = clip_loss + caption_loss
        total_loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    # Epoch summary
    print(f"Epoch {epoch}:")
    print(f"  Avg CLIP loss: {sum(metrics['clip_loss'])/len(metrics['clip_loss']):.4f}")
    print(f"  Avg Caption loss: {sum(metrics['caption_loss'])/len(metrics['caption_loss']):.4f}")
```

## Mathematical Formulation

The total CoCa loss is:

$\mathcal{L}_{\text{total}} = w_1 \cdot \mathcal{L}_{\text{contrastive}} + w_2 \cdot \mathcal{L}_{\text{caption}}$

Where:

1. **Contrastive Loss** (inherited from ClipLoss):
   $\mathcal{L}_{\text{contrastive}} = \frac{1}{2}\left[\text{CE}(\tau \cdot I T^\top, y) + \text{CE}(\tau \cdot T I^\top, y)\right]$

2. **Caption Loss** (cross-entropy with teacher forcing):
   $\mathcal{L}_{\text{caption}} = -\frac{1}{N} \sum_{i=1}^{N} \sum_{t=1}^{T} \mathbb{1}[y_i^t \neq \text{pad}] \cdot \log P(y_i^t | y_i^{<t}, x_i)$

   Where:

   * $N$ = batch size
   * $T$ = sequence length
   * $y_i^t$ = target token at position $t$
   * $x_i$ = image features
   * Padding tokens are ignored via `ignore_index`

## Hyperparameter Tuning

**Recommended loss weight ratios:**

| Dataset Type           | Caption Weight | CLIP Weight | Notes                         |
| ---------------------- | -------------- | ----------- | ----------------------------- |
| Large web data (LAION) | 1.0-2.0        | 1.0         | Balanced training             |
| Caption-focused (COCO) | 2.0-3.0        | 1.0         | Prioritize generation quality |
| Retrieval-focused      | 0.5-1.0        | 1.0         | Prioritize matching           |
| Fine-tuning            | 3.0-5.0        | 0.1-0.5     | Adapt caption style           |

**Guidelines:**

* Start with `caption_loss_weight=2.0`, `clip_loss_weight=1.0`
* If captions are low quality, increase `caption_loss_weight`
* If retrieval performance is poor, increase `clip_loss_weight`
* Monitor both loss components separately

## Related

* [CoCa](/api/coca-model) - Model that uses this loss
* [ClipLoss](/api/clip-loss) - Parent class for contrastive loss
* [Training Guide](/training/overview) - Full training examples
